"""
数据管理器 - 提供统一的数据访问接口
集成向量存储、图谱存储和文档存储
"""

from typing import List, Dict, Any, Optional
import asyncio

from ..core.interfaces import (
    VectorStoreInterface, GraphStoreInterface, DocumentStoreInterface,
    Document, SearchResult, Entity, Relationship
)
from ..core.events import EventEmitter
from ..core.container import container


class DataManager(EventEmitter):
    """
    统一数据管理器
    协调各种存储后端的操作
    """
    
    def __init__(self):
        super().__init__()
        self.vector_store: Optional[VectorStoreInterface] = None
        self.graph_store: Optional[GraphStoreInterface] = None
        self.document_store: Optional[DocumentStoreInterface] = None
        self._initialized = False
    
    async def initialize(self):
        """初始化所有存储后端"""
        if self._initialized:
            return
        
        try:
            # 从容器解析存储实例
            self.vector_store = await container.resolve(VectorStoreInterface)
            self.graph_store = await container.resolve(GraphStoreInterface)
            self.document_store = await container.resolve(DocumentStoreInterface)
            
            # 初始化各个存储
            await self.vector_store.initialize()
            await self.graph_store.initialize()
            
            self._initialized = True
            
            await self.emit("data_manager_initialized", {
                "vector_store": type(self.vector_store).__name__,
                "graph_store": type(self.graph_store).__name__,
                "document_store": type(self.document_store).__name__
            })
            
        except Exception as e:
            await self.emit_error("data_manager_initialization", e)
            raise
    
    async def store_document_complete(
        self, 
        doc_id: str, 
        content: str, 
        raw_content: bytes,
        metadata: Dict[str, Any],
        embedding: Optional[List[float]] = None
    ) -> bool:
        """
        完整存储文档到所有后端
        
        Args:
            doc_id: 文档ID
            content: 文档文本内容
            raw_content: 原始文档内容
            metadata: 文档元数据
            embedding: 文档嵌入向量
        """
        if not self._initialized:
            await self.initialize()
        
        try:
            # 并行执行所有存储操作
            tasks = []
            
            # 1. 存储原始文档
            tasks.append(
                self.document_store.store_document(doc_id, raw_content, metadata)
            )
            
            # 2. 存储到向量数据库
            document = Document(
                id=doc_id,
                content=content,
                metadata=metadata,
                embedding=embedding
            )
            tasks.append(
                self.vector_store.add_documents([document])
            )
            
            # 3. 如果有实体信息，存储到图数据库
            entities = metadata.get("entities", [])
            relationships = metadata.get("relationships", [])
            
            if entities:
                for entity_data in entities:
                    entity = Entity(**entity_data)
                    tasks.append(self.graph_store.add_entity(entity))
            
            if relationships:
                for rel_data in relationships:
                    relationship = Relationship(**rel_data)
                    tasks.append(self.graph_store.add_relationship(relationship))
            
            # 等待所有任务完成
            results = await asyncio.gather(*tasks, return_exceptions=True)
            
            # 检查是否有失败的操作
            failures = [r for r in results if isinstance(r, Exception)]
            if failures:
                await self.emit("document_storage_partial_failure", {
                    "document_id": doc_id,
                    "failures": [str(f) for f in failures]
                })
                return False
            
            await self.emit("document_stored_complete", {
                "document_id": doc_id,
                "metadata": metadata
            })
            
            return True
            
        except Exception as e:
            await self.emit_error("store_document_complete", e)
            return False
    
    async def hybrid_search(
        self,
        query: str,
        query_embedding: Optional[List[float]] = None,
        top_k: int = 10,
        search_modes: List[str] = ["vector", "graph"],
        metadata_filter: Optional[Dict[str, Any]] = None
    ) -> List[SearchResult]:
        """
        混合搜索 - 结合向量搜索和图搜索
        
        Args:
            query: 查询文本
            query_embedding: 查询嵌入向量
            top_k: 返回结果数量
            search_modes: 搜索模式列表
            metadata_filter: 元数据过滤条件
        """
        if not self._initialized:
            await self.initialize()
        
        try:
            all_results = []
            
            # 向量搜索
            if "vector" in search_modes:
                if query_embedding:
                    vector_results = await self.vector_store.search_similar(
                        query_embedding, top_k, metadata_filter
                    )
                else:
                    vector_results = await self.vector_store.search_by_text(
                        query, top_k, metadata_filter
                    )
                all_results.extend(vector_results)
            
            # 图搜索（基于实体）
            if "graph" in search_modes and self.graph_store:
                graph_results = await self._graph_based_search(query, top_k)
                all_results.extend(graph_results)
            
            # 结果去重和排序
            unique_results = self._deduplicate_results(all_results)
            sorted_results = sorted(unique_results, key=lambda x: x.score, reverse=True)
            
            await self.emit("hybrid_search_completed", {
                "query": query,
                "search_modes": search_modes,
                "total_results": len(sorted_results),
                "top_k": top_k
            })
            
            return sorted_results[:top_k]
            
        except Exception as e:
            await self.emit_error("hybrid_search", e)
            return []
    
    async def _graph_based_search(self, query: str, top_k: int) -> List[SearchResult]:
        """基于图谱的搜索"""
        try:
            # 简单的实体匹配搜索
            # 在实际应用中，这里可以使用更复杂的图查询
            cypher_query = """
            MATCH (n)
            WHERE n.name CONTAINS $query OR n.title CONTAINS $query
            RETURN n, labels(n) as types
            LIMIT $limit
            """
            
            graph_results = await self.graph_store.query_graph(
                cypher_query, 
                {"query": query, "limit": top_k}
            )
            
            search_results = []
            for record in graph_results:
                node = record["n"]
                node_types = record["types"]
                
                # 转换为SearchResult
                document = Document(
                    id=node.get("id", ""),
                    content=f"{node.get('name', '')} - {node.get('description', '')}",
                    metadata={
                        "source": "graph",
                        "entity_type": node_types[0] if node_types else "Unknown",
                        "properties": dict(node)
                    }
                )
                
                search_result = SearchResult(
                    document=document,
                    score=0.8,  # 图搜索的默认分数
                    relevance_type="graph"
                )
                search_results.append(search_result)
            
            return search_results
            
        except Exception as e:
            await self.emit_error("graph_based_search", e)
            return []
    
    def _deduplicate_results(self, results: List[SearchResult]) -> List[SearchResult]:
        """结果去重"""
        seen_ids = set()
        unique_results = []
        
        for result in results:
            if result.document.id not in seen_ids:
                seen_ids.add(result.document.id)
                unique_results.append(result)
        
        return unique_results
    
    async def get_entity_context(self, entity_id: str, max_depth: int = 2) -> Dict[str, Any]:
        """获取实体的上下文信息"""
        if not self._initialized:
            await self.initialize()
        
        try:
            # 获取相关实体
            related_entities = await self.graph_store.find_related_entities(
                entity_id, max_depth=max_depth
            )
            
            # 获取邻居信息
            neighbors = await self.graph_store.get_entity_neighbors(entity_id)
            
            # 构建上下文
            context = {
                "entity_id": entity_id,
                "related_entities": [
                    {
                        "id": entity.id,
                        "type": entity.type,
                        "properties": entity.properties
                    }
                    for entity in related_entities
                ],
                "neighbors_by_relation": {}
            }
            
            # 整理邻居信息
            for relation_type, entities in neighbors.items():
                context["neighbors_by_relation"][relation_type] = [
                    {
                        "id": entity.id,
                        "type": entity.type,
                        "properties": entity.properties
                    }
                    for entity in entities
                ]
            
            return context
            
        except Exception as e:
            await self.emit_error("get_entity_context", e)
            return {}
    
    async def get_system_stats(self) -> Dict[str, Any]:
        """获取整个系统的统计信息"""
        if not self._initialized:
            await self.initialize()
        
        try:
            # 并行获取各个存储的统计信息
            tasks = []
            
            if hasattr(self.vector_store, 'get_collection_stats'):
                tasks.append(self.vector_store.get_collection_stats())
            else:
                tasks.append(asyncio.create_task(asyncio.sleep(0, result={})))
            
            if hasattr(self.graph_store, 'get_graph_stats'):
                tasks.append(self.graph_store.get_graph_stats())
            else:
                tasks.append(asyncio.create_task(asyncio.sleep(0, result={})))
            
            if hasattr(self.document_store, 'get_storage_stats'):
                tasks.append(self.document_store.get_storage_stats())
            else:
                tasks.append(asyncio.create_task(asyncio.sleep(0, result={})))
            
            vector_stats, graph_stats, document_stats = await asyncio.gather(*tasks)
            
            return {
                "vector_store": vector_stats,
                "graph_store": graph_stats,
                "document_store": document_stats,
                "timestamp": asyncio.get_event_loop().time()
            }
            
        except Exception as e:
            await self.emit_error("get_system_stats", e)
            return {}
    
    async def cleanup(self):
        """清理资源"""
        try:
            if self.graph_store and hasattr(self.graph_store, 'close'):
                await self.graph_store.close()
            
            await self.emit("data_manager_cleanup_completed", {})
            
        except Exception as e:
            await self.emit_error("data_manager_cleanup", e)